import numpy as np
import numpy
import pandas as pd


def food_map(series):
    print(type(series))
    if series['food'] == 'A1':
        return 'A'
    elif series['food'] == 'A2':
        return 'A'
    elif series['food'] == 'B1':
        return 'B'
    elif series['food'] == 'B2':
        return 'B'
    elif series['food'] == 'B3':
        return 'B'
    elif series['food'] == 'C1':
        return 'C'
    elif series['food'] == 'C2':
        return 'C'


if __name__ == '__main__':
    df1: pd.DataFrame = pd.DataFrame(
        {'food': ['A1', 'A2', 'B1', 'B2', 'B3', 'C1', 'C2'], 'data': [1, 2, 3, 4, 5, 6, 7]})
    # 需要说明，food_map是一个Series  是一行数据 {'food':'A1','data':1}这样的
    df1['food_map'] = df1.apply(food_map, axis='columns')
